import grpc
import p4runtime_pb2
import p4tuntime_pb2_grpc
from p4.v1 import p4runtime_pb2
from p4.v1 import p4runtime_pb2_grpc
from p4runtime_lib.helper import P4InfoHelper
form p4runtime_lib.bmv2 import Bmv2SwitchConnection
import time

def set_ecn_threshold(switch, p4info_helper, port, threshold):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyEgress.ecn_threshold_table",
        match_fields={
            "standards_metadata.egress_port": port
        },
        action_name="MyEgress.set_ecn_threshold",
        action_params={
            "threshold": threshold
        }
    )

    write_request = p4runtime_pb2.WriteRequest()
    write_request.device_id = switch.device_id
    update = write_request.updates.add()
    update.type = p4runtime_pb2.Update.INSERT
    update.entity.table_entry.CopyFrom(table_entry)
    
def main();
    p4info_file_path = "build/ecn.p4.p4info.txtpb"
    bmv2_file_path = "build/ecn.json"

    p4info_helper = P4InfoHelper(p4info_file_path)

    s1 = Bmv2SwitchConnection(
        name="s1",
        address="127.0.0.1:50051",
        device_id=0,
        proto_dump_file='logs/s1-p4runtime-request.txt'
    )

    s1 = Bmv2SwitchConnection(
        name="s1",
        address="127.0.0.1:50052",
        device_id=1,
        proto_dump_file='logs/s2-p4runtime-request.txt'
    )

    s1 = Bmv2SwitchConnection(
        name="s1",
        address="127.0.0.1:50053",
        device_id=2,
        proto_dump_file='logs/s3-p4runtime-request.txt'
    )

    s1.MasterArbitrationUpdate()
    s2.MasterArbitrationUpdate()
    s3.MasterArbitrationUpdate()

    s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,bmv2_json_file_path=bmv2_file_path)\
    print("Installed P4 Program using SetForwardingPipelineConfig on s1")

    s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,bmv2_json_file_path=bmv2_file_path)\
    print("Installed P4 Program using SetForwardingPipelineConfig on s2")

    s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,bmv2_json_file_path=bmv2_file_path)\
    print("Installed P4 Program using SetForwardingPipelineConfig on s3")

    set_ecn_threshold(s1,p4info_helper, 1, 100)
    set_ecn_threshold(s2,p4info_helper, 1, 100)
    set_ecn_threshold(s3,p4info_helper, 1, 100)

    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print(" Shutting down")
    except grpc.RpcError as e:
        print(f"gRPC error:{e} ")

    s1.shutdown()
    s2.shutdown()
    s3.shutdown()

if__name__ == "__main__":
    main()
